{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "## Training a Variational Quantum Eigensolver with evotorch and pennylane\n",
    "\n",
    "This example demonstrates how you can train variational quantum eigensolvers (VQEs) using EvoTorch and PennyLane. To execute this example, you will need to install PennyLane with:\n",
    "\n",
    "```\n",
    "    pip install pennylane\n",
    "```\n",
    "\n",
    "This example is based on PennyLane's [VQE example](https://pennylane.ai/qml/demos/tutorial_vqe.html) and we encourage the reader to read that tutorial for a deeper understanding of the problem setting. \n",
    "\n",
    "### Basics of variational quantum eigensolvers\n",
    "\n",
    "VQE is a recently introduced algorithm which has the objective to train a quantum computer to prepare the ground state of a given molecule [1]. Finding this ground state is a central problem in computational physics and quantum chemistry, as knowledge of this state then enables parameterisation of further simulations. \n",
    "\n",
    "In brief, a VQE consists of 3 components:\n",
    "\n",
    "1. A parameterisable quantum circuit `Q` with `p` parameters, which prepares the ground state of the molecule.\n",
    "2. A cost function `C` that computes the energy of a given ground state, which we want to minimise. \n",
    "3. A classical optimisation algorithm, which searches for the optimal parameter vector which minimises the energy.\n",
    "\n",
    "Here we will be using variation quantum circuits for `Q`, which means that the circuits are parameterised by a vector of angles. \n",
    "\n",
    "### Limitations of backpropagation\n",
    "\n",
    "The most natural solution to this problem would be to use backpropagation to train the parameters of the quantum circuit to minimise the cost. However, this is only possible in simulation, as backpropagation through a physical quantum circuit would require observing and reusing the state values of the circuit, which is impossible. Because of this, alternative approaches to differentiation of quantum circuits have been proposed. A particularly prominent approach is the 'parameter-shift rule' [2] which, effectively allows us to compute analytic gradients on the parameters by evaluating precise shifts of the parameters, one-by-one. This, of course, means that for one update of a circuit with `p` parameters, we must perform `2p` circuit evaluations.\n",
    "\n",
    "A recent paper [3] introduced the idea that evolutionary algorithms such as SNES and XNES may serve as an alternative approach to quantum circuit optimisation, where the number of circuit evaluation in each iteration is simply the population size `k`. In practice, this can yield similar performance to gradient-descent based methods in less circuit evaluations. In this example, we'll be doing exactly that: training the parameters of a VQE using SNES."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "### Setting up the Cost Function\n",
    "\n",
    "First, we will need to define the molecular structure. We will study the H$_2$ molecule, which means that we have:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "symbols = [\"H\", \"H\"]  #H2 molecule\n",
    "coordinates = torch.tensor([0.0, 0.0, -0.6614, 0.0, 0.0, 0.6614])  # nuclear coordinates in atomic units"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "Next, we need to calculate the electronic Hamiltonian:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pennylane as qml\n",
    "\n",
    "H, qubits = qml.qchem.molecular_hamiltonian(symbols, coordinates)\n",
    "print(\"Number of qubits = \", qubits)\n",
    "print(\"The Hamiltonian is \", H)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "We'll be using the default PennyLane quantum simulator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "dev = qml.device(\"default.qubit\", wires=qubits)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "Now we're ready to set up the circuit `Q` and cost function `C`. Note that due to the simplicity of the molecule, we have only a single parameter, `p=1` to optimise!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "electrons = 2  # H2 has 2 electrons\n",
    "hf = qml.qchem.hf_state(electrons, qubits)  # giving the Hartree-Fock state \n",
    "\n",
    "def circuit(param, wires):\n",
    "    # Setting up the circuit to optimise, which simply consists of preparing the basis state as the Hartree-Fock state \n",
    "    # And then applying a double excitation parameterised by the parameter\n",
    "    qml.BasisState(hf, wires=wires)\n",
    "    qml.DoubleExcitation(param, wires=[0, 1, 2, 3])\n",
    "    \n",
    "@qml.qnode(dev, diff_method=None)  # Disabling gradients -- we don't need them\n",
    "def cost_fn(param):\n",
    "    # Defining the cost function: simply apply the parameterised circuit and take the exponent of the Hamiltonian\n",
    "    circuit(param, wires=range(qubits))\n",
    "    return qml.expval(H)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "### Creating a EvoTorch Problem instance\n",
    "\n",
    "With the cost function well-defined, we're ready to create a `Problem` instance. Note that we will repeat the steps above, so that the model has its own internal definition of the cost function. This will allow us to exploit parallelization with Ray. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "from evotorch import Problem, Solution\n",
    "from evotorch.algorithms import SNES\n",
    "from evotorch.logging import StdOutLogger, PandasLogger\n",
    "\n",
    "from typing import Optional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VGEMin(Problem):\n",
    "    \n",
    "    def __init__(self, num_actors: Optional[int] = None):\n",
    "        \n",
    "        super().__init__(\n",
    "            objective_sense='min',  # Minimise the objective\n",
    "            solution_length = 1,  # There is only a single parameter to optimise\n",
    "            initial_bounds = (-1e-6, 1e-6),  # Start the search very close to zero\n",
    "            num_actors = num_actors,  # Number of ray actors\n",
    "        )\n",
    "        \n",
    "        symbols = [\"H\", \"H\"]\n",
    "        coordinates = torch.tensor([0.0, 0.0, -0.6614, 0.0, 0.0, 0.6614])\n",
    "        self._H, self._qubits = qml.qchem.molecular_hamiltonian(symbols, coordinates)\n",
    "        \n",
    "        electrons = 2  # H2 has 2 electrons\n",
    "        self._hf = qml.qchem.hf_state(electrons, qubits)\n",
    "        \n",
    "    def _prepare(self):\n",
    "        # Prepare function called by actors once instantaited\n",
    "        dev = qml.device(\"default.qubit\", wires=self._qubits)\n",
    "        \n",
    "        # Inline definition of cost function allows us to easily decorate it as a quantum node\n",
    "        @qml.qnode(dev, diff_method = None, interface = 'torch')\n",
    "        def actor_cost_fn(param):\n",
    "            with torch.no_grad():\n",
    "                wires = range(self._qubits)\n",
    "                qml.BasisState(self._hf, wires=wires)\n",
    "                qml.DoubleExcitation(param[0], wires=[0, 1, 2, 3])\n",
    "                return qml.expval(self._H)\n",
    "            \n",
    "        self._cost_fn = actor_cost_fn\n",
    "        \n",
    "    def _evaluate(self, individual: Solution):\n",
    "        x = individual.access_values()  # Get the decision values -- in this case a vector of length 1\n",
    "        cost = self._cost_fn(x)  # Evaluate the decision values\n",
    "        individual.set_evals(cost)  # Update the fitness\n",
    "        \n",
    "problem = VGEMin(num_actors = 4)  # Instantiate the problem class\n",
    "population = problem.generate_batch(5)  # Generate a population to test things out\n",
    "problem.evaluate(population)  # If we've set everything up correctly we should get no error\n",
    "print(f'Initial fitness values {population.access_evals()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "### Training the VQE with EvoTorch\n",
    "\n",
    "Now we're ready to train. Simply create the searcher and logger:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "searcher = SNES(problem, stdev_init=0.1)  # stdev_init=0.1 used in [3]\n",
    "logger = PandasLogger(searcher)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "And train for 100 generations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "searcher.run(100)\n",
    "\n",
    "progress = logger.to_dataframe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "progress.mean_eval.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "Taking a look at the mean of the searcher we have:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Final mean is {searcher.status[\"center\"]} Final stdev is {searcher.status[\"stdev\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "Where the learned mean is close to the known approximate global optima of 0.208.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Cost of learned mean is {cost_fn(searcher.status[\"center\"][0].numpy())} vs. approx global optima {cost_fn(0.208)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "### A more challenging example\n",
    "\n",
    "With the basics down, we can now think about a more challenging problem. Instead of the simple H$_2$ molecule, we will instead consider the water molecule H2O. A closely related experiment was performed in [3], meaning that this is a demonstration of a relatively 'state of the art' result for QVE training using EvoTorch.\n",
    "\n",
    "Now we have the molecule configuration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "symbols = [\"H\", \"O\", \"H\"]  #H2O molecule\n",
    "coordinates = torch.tensor([-0.0399, -0.0038, 0.0, 1.5780, 0.8540, 0.0, 2.7909, -0.5159, 0.0])  # nuclear coordinates in atomic units"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "And as before we need the electronic Hamiltonian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "H, qubits = qml.qchem.molecular_hamiltonian(\n",
    "    symbols, \n",
    "    coordinates,\n",
    "    charge=0,\n",
    "    mult=1,\n",
    "    basis=\"sto-3g\",\n",
    "    active_electrons=4,\n",
    "    active_orbitals=4,\n",
    ") \n",
    "\n",
    "print(\"Number of qubits = \", qubits)\n",
    "print(\"The Hamiltonian is \", H)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": [
    "Making a new cost function, this time using UCCSD ansatz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "electrons = 10\n",
    "orbitals = 7\n",
    "core, active = qml.qchem.active_space(electrons, orbitals, active_electrons=4, active_orbitals=4)\n",
    "\n",
    "singles, doubles = qml.qchem.excitations(len(active), qubits)\n",
    "hf = qml.qchem.hf_state(\n",
    "    len(active), \n",
    "    qubits,\n",
    ")  # giving the Hartree-Fock state \n",
    "\n",
    "# Map excitations to the wires the UCCSD circuit will act on\n",
    "s_wires, d_wires = qml.qchem.excitations_to_wires(singles, doubles)\n",
    "\n",
    "# Define the device\n",
    "dev = qml.device('default.qubit', wires=qubits)\n",
    "\n",
    "def circuit2(param, wires):\n",
    "    # Setting up the circuit to optimise, which simply consists of preparing the basis state as the Hartree-Fock state \n",
    "    # And then applying a UCCSD ansatz\n",
    "    qml.UCCSD(param, wires=wires, s_wires = s_wires, d_wires = d_wires, init_state = hf)\n",
    "    \n",
    "@qml.qnode(dev, diff_method=None)  # Disabling gradients -- we don't need them\n",
    "def cost_fn(param):\n",
    "    # Defining the cost function: simply apply the parameterised circuit and take the exponent of the Hamiltonian\n",
    "    circuit2(param, wires=range(qubits))\n",
    "    return qml.expval(H)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27",
   "metadata": {},
   "source": [
    "Putting this together in a problem definition:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining a new problem:\n",
    "class VGEH2O(Problem):\n",
    "    \n",
    "    def __init__(self, num_actors: Optional[int]):\n",
    "        \n",
    "        super().__init__(\n",
    "            objective_sense='min',  # Minimise the objective\n",
    "            solution_length = 26,  # There are 26 parameters to optimise\n",
    "            initial_bounds = (-1e-6, 1e-6),  # Start the search very close to zero\n",
    "            num_actors = num_actors,\n",
    "        )\n",
    "        \n",
    "    def _prepare(self):\n",
    "        \n",
    "        symbols = [\"H\", \"O\", \"H\"]  #H2O molecule\n",
    "        coordinates = torch.tensor([-0.0399, -0.0038, 0.0, 1.5780, 0.8540, 0.0, 2.7909, -0.5159, 0.0])\n",
    "        self._H, self._qubits = qml.qchem.molecular_hamiltonian(\n",
    "            symbols, \n",
    "            coordinates,\n",
    "            charge=0,\n",
    "            mult=1,\n",
    "            basis=\"sto-3g\",\n",
    "            active_electrons=4,\n",
    "            active_orbitals=4,\n",
    "        ) \n",
    "\n",
    "        electrons = 10\n",
    "        orbitals = 7\n",
    "        core, active = qml.qchem.active_space(electrons, orbitals, active_electrons=4, active_orbitals=4)\n",
    "\n",
    "        singles, doubles = qml.qchem.excitations(len(active), qubits)\n",
    "        self._hf = qml.qchem.hf_state(\n",
    "            len(active), \n",
    "            qubits,\n",
    "        ) \n",
    "\n",
    "        self._s_wires, self._d_wires = qml.qchem.excitations_to_wires(singles, doubles)\n",
    "        \n",
    "        # Prepare function called by actors once instantaited\n",
    "        dev = qml.device(\"default.qubit\", wires=self._qubits)\n",
    "        \n",
    "        # Inline definition of cost function allows us to easily decorate it as a quantum node\n",
    "        @qml.qnode(dev, diff_method = None, interface = 'torch')\n",
    "        def actor_cost_fn(param):\n",
    "            with torch.no_grad():\n",
    "                wires = range(self._qubits)\n",
    "                qml.UCCSD(param, wires=wires, s_wires = self._s_wires, d_wires = self._d_wires, init_state = self._hf)\n",
    "                return qml.expval(self._H)\n",
    "            \n",
    "        self._cost_fn = actor_cost_fn\n",
    "        \n",
    "    def _evaluate(self, individual: Solution):\n",
    "        x = individual.access_values()  # Get the decision values\n",
    "        cost = self._cost_fn(x)  # Evaluate the parameter vector\n",
    "        individual.set_evals(cost)  # Update the fitness\n",
    "        \n",
    "problem = VGEH2O(num_actors = 4)\n",
    "population = problem.generate_batch(10)\n",
    "problem.evaluate(population)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29",
   "metadata": {},
   "source": [
    "And from there, all we need to do is train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "searcher = SNES(problem, stdev_init=0.1)  # stdev_init=0.1 used in [3]\n",
    "pandas_logger = PandasLogger(searcher)\n",
    "stdout_logger = StdOutLogger(searcher)\n",
    "\n",
    "# Run for 200 generations\n",
    "searcher.run(200)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31",
   "metadata": {},
   "source": [
    "And visualize the progress:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "progress = pandas_logger.to_dataframe()\n",
    "progress.mean_eval.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33",
   "metadata": {},
   "source": [
    "#### References\n",
    "\n",
    "[1] Peruzzo, Alberto, et al. [\"A variational eigenvalue solver on a photonic quantum processor.\"](https://www.nature.com/articles/ncomms5213) Nature communications 5.1 (2014): 1-7.\n",
    "\n",
    "[2] Schuld, Maria, et al. [\"Evaluating analytic gradients on quantum hardware.\"](https://journals.aps.org/pra/abstract/10.1103/PhysRevA.99.032331) Physical Review A 99.3 (2019): 032331.\n",
    "\n",
    "[3] Anand, Abhinav, Matthias Degroote, and Alán Aspuru-Guzik. [\"Natural evolutionary strategies for variational quantum computation.\"](https://iopscience.iop.org/article/10.1088/2632-2153/abf3ac). Machine Learning: Science and Technology 2.4 (2021): 045012"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
